// Copyright (c) 2025 Huawei Technologies Co., Ltd.
// openUBMC is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.
//
// Description: KVM独立客户端接口定义

#include <cstdlib>
#include <iostream>
#include <thread>
#include <memory>
#include <string>
#include "KvmClient.hpp"

KvmClient::KvmClient() {
    RegisterRecvMessageHandler(CMD_CONNECT_RSP_S, [this](const std::vector<uint8_t>& body) {
        ProcessConnectRsp(body);
    });
}

KvmClient::~KvmClient() {
    DeleteCharArray(token_);
    StopHeartbeat();
    Disconnect();
}

// 注册错误处理函数
void KvmClient::SetErrorHandler(std::function<void(const std::exception&)> handler) {
    errorHandler_ = std::move(handler);
    vmmClient_.SetErrorHandler(errorHandler_);
}

// 注册会话关闭处理函数
void KvmClient::SetCloseHandler(std::function<void(const uint16_t&)> handler) {
    closeHandler_ = std::move(handler);
    vmmClient_.SetCloseHandler(closeHandler_);
}

// 注册日志处理函数
void KvmClient::SetLogHandler(std::function<void(const std::string&)> handler) {
    logHandler_ = std::move(handler);
    vmmClient_.SetLogHandler(logHandler_);
}

void KvmClient::RegisterSendMessageHandler(
    uint16_t cmdType, std::function<std::pair<int, std::vector<uint8_t>>(const std::vector<uint8_t>&)> handler) {
    sendMessageHandlers_[cmdType] = std::move(handler);
}

void KvmClient::RegisterRecvMessageHandler(
    uint16_t cmdType, std::function<void(const std::vector<uint8_t>&)> handler) {
    recvMessageHandlers_[cmdType] = std::move(handler);
}

int KvmClient::CreateSession(const char* host, const char* port,
    const char* userName, const char* password, const char* kvmMode) {
    if (!authManager_) {
        authManager_ = std::make_unique<AuthenticationManager>();
    }

    auto create_res = authManager_->HttpCreateSession(host, port, userName, password, kvmMode);
    if (!create_res) {
        LogMessage("OpenKvm: Failed to create session.");
        return E_ERR;
    }

    auto info = authManager_->GetCreateSessionResultInfo();
    const HttpSuccInfo* pSuccType = boost::get<HttpSuccInfo>(&info);
    const HttpFailInfo* pFailType = boost::get<HttpFailInfo>(&info);

    if (!pSuccType) {
        if (!pFailType) {
            LogMessage("OpenKvm: Unknown create session result type.");
            return E_ERR;
        }

        auto it = ERROR_CODE_MAP.find(pFailType->errorCode);
        if (it != ERROR_CODE_MAP.end()) {
            return it->second; // 返回对应的错误码
        }

        // 如果错误码未在映射中定义，返回通用错误
        return E_ERR;
    }

    return E_OK;
}

int KvmClient::OpenKvm(const char* host, const char* port,
    const char* userName, const char* password, const char* kvmMode) {
    
    // 防止重入
    std::lock_guard<std::mutex> lock(openKvmMutex_);

    if (isConnected_) {
        LogMessage("OpenKvm: Already connected to KVM.");
        return E_ALREADY_CONNECTED;
    }

    try {
        // 创建会话
        int sessionResult = CreateSession(host, port, userName, password, kvmMode);
        if (sessionResult != E_OK) {
            return sessionResult;
        }

        auto info = authManager_->GetCreateSessionResultInfo();
        const HttpSuccInfo* pSuccType = boost::get<HttpSuccInfo>(&info);
        Connect(host, std::to_string(pSuccType->kvmPort), "/websocket/kvm");
        std::thread t([this]() { Run(); });
        std::this_thread::sleep_for(std::chrono::seconds(1));   // 稍微等待1秒保证I/O服务已启动
        
        PackData pack_data;
        // 初始化连接响应变量
        connectRsp_ = INIT_RESPONSE_CODE;

        Send(pack_data.ConnectKvm(pSuccType->token));

        // 发送连接请求后，等待服务端的响应
        std::unique_lock<std::mutex> rspLock(connectRspMutex_);
        if (!connectRspCv_.wait_for(rspLock, std::chrono::seconds(RESPONSE_TIMEOUT),
            [this]() { return connectRsp_ != INIT_RESPONSE_CODE; })) {
            LogMessage("OpenKvm: Timeout waiting for connection response.");
            if (authManager_) {
                authManager_->HttpDeleteSession();
            }
            return E_TIMEOUT;
        }

        StartHeartbeat(std::chrono::seconds(HEARTBEAT_INTERVAL));
        t.detach();

        // 返回服务端的结果
        if (connectRsp_ == E_OK) {
            isConnected_ = true;
            token_ = pSuccType->token;
            host_ = host;
            LogMessage("Open KVM successfully.");
        } else {
            LogMessage("Open KVM failed.");
            Disconnect();
        }
        return connectRsp_;
    }
    catch(const std::exception& e) {
        LogMessage(std::string("Exception in OpenKvm: ") + e.what());
        Disconnect();
        OnError(e);
        return E_ERR;
    }
}

int KvmClient::CloseKvm() {
    if (!isConnected_) {
        LogMessage("Not connected to KVM.");
        return E_NOT_CONNECTED;
    }

    try {
        PackData pack_data;
        Send(pack_data.DisconnectKvm());
        StopHeartbeat();
    } catch (const std::exception& e) {
        LogMessage(std::string("Exception in CloseKvm: ") + e.what());
        if (errorHandler_) {
            errorHandler_(e);
        }
        return E_ERR;
    }
    
    isConnected_ = false;
    LogMessage("Close KVM successfully.");
    return E_OK;
}

int KvmClient::OpenVmm(const std::string& filePath, MountType type) {
    if (!isConnected_) {
        LogMessage("OpenVmm: Not connected to KVM.");
        return E_NOT_CONNECTED;
    }

    if (type != MountType::IsoImage) {
        LogMessage("OpenVmm: Invalid mount type.");
        return E_INVALID_PARAM;
    }

    FileType fileType = GetFileType(filePath);
    if (fileType == FileType::INVALID || static_cast<int>(fileType) != static_cast<int>(type)) {
        // 如果文件类型不匹配，返回错误
        LogMessage("OpenVmm: Invalid file path.");
        return E_INVALID_FILE_PATH;
    }

    // 获取VMM使能是否开启
    auto itGetStatus = sendMessageHandlers_.find(CMD_GET_VMM_STATUS_REQ_C);
    if (itGetStatus == sendMessageHandlers_.end()) {
        LogMessage("OpenVmm: No handler registered.");
        return E_ERR;
    }

    // 调用注册的回调函数
    std::pair<int, std::vector<uint8_t>> getStatusRes = (itGetStatus->second({}));
    if (getStatusRes.first != E_OK) {
        LogMessage("OpenVmm: Failed to get VMM status.");
        return getStatusRes.first;
    }
    uint8_t vmmStatus = getStatusRes.second[0];
    if (vmmStatus != 1) {
        LogMessage("VMM is disabled.");
        return E_VMM_DISABLED;
    }

    // 获取VMM端口号
    auto itGetPort = sendMessageHandlers_.find(CMD_GET_VMM_PORT_REQ_C);
    if (itGetPort == sendMessageHandlers_.end()) {
        LogMessage("OpenVmm: No handler registered.");
        return E_ERR;
    }

    // 调用注册的回调函数
    std::pair<int, std::vector<uint8_t>> getPortRes = (itGetPort->second({}));
    if (getPortRes.first != E_OK) {
        LogMessage("OpenVmm: Failed to get VMM port.");
        return getPortRes.first;
    }
    // 2字节的端口号
    uint16_t vmmPort = static_cast<uint16_t>(getPortRes.second[0] << 8) + static_cast<uint16_t>(getPortRes.second[1]);

    return vmmClient_.OpenVmm(filePath, type, host_, vmmPort, token_);
}

int KvmClient::CloseVmm() {
    if (!isConnected_) {
        LogMessage("CloseVmm: Not connected to KVM.");
        return E_NOT_CONNECTED;
    }

    return vmmClient_.CloseVmm();
}

int KvmClient::GetBootOption(BootOption& option)
{
    if (!isConnected_) {
        LogMessage("Get boot option failed, for not connected to KVM.");
        return E_NOT_CONNECTED;
    }

    auto it = sendMessageHandlers_.find(CMD_GET_BOOT_OPTION_REQ_C);
    if (it == sendMessageHandlers_.end()) {
        LogMessage("GetBootOption: No handler registered.");
        return E_ERR;
    }

    // 调用注册的回调函数
    PackData pack_data;
    uint8_t op = static_cast<uint8_t>(option);
    std::vector<uint8_t> body = {op};
    std::pair<int, std::vector<uint8_t>> result = (it->second(body));
    option = static_cast<BootOption>(result.second[0]);
    return result.first;
}

int KvmClient::SetBootOption(BootOption option)
{
    if (!isConnected_) {
        LogMessage("Set boot option failed, for not connected to KVM.");
        return E_NOT_CONNECTED;
    }

    auto it = sendMessageHandlers_.find(CMD_SET_BOOT_OPTION_REQ_C);
    if (it == sendMessageHandlers_.end()) {
        LogMessage("SetBootOption: No handler registered.");
        return E_ERR;
    }

    // 调用注册的回调函数
    PackData pack_data;
    uint8_t op = static_cast<uint8_t>(option);
    std::vector<uint8_t> body = {op};
    return (it->second(body)).first;
}

int KvmClient::PowerControl(PowerControlOption option)
{
    if (!isConnected_) {
        LogMessage("Power control failed, for not connected to KVM.");
        return E_NOT_CONNECTED;
    }

    auto it = sendMessageHandlers_.find(CMD_POWER_CONTROL_REQ_C);
    if (it == sendMessageHandlers_.end()) {
        LogMessage("PowerControl: No handler registered.");
        return E_ERR;
    }

    // 调用注册的回调函数
    PackData pack_data;
    uint8_t op = static_cast<uint8_t>(option);
    std::vector<uint8_t> body = {op};
    return (it->second(body)).first;
}

void KvmClient::OnClose(const websocket::close_reason& cr) {
    isConnected_ = false;

    if (closeHandler_) {
        closeHandler_(cr.code);
    }
}

void KvmClient::OnError(const std::exception& error) {
    isConnected_ = false;

    if (authManager_) {
        authManager_->HttpDeleteSession();
    }
    if (errorHandler_) {
        errorHandler_(error);
    }
}

void KvmClient::LogMessage(const std::string& message) {
    if (logHandler_) {
        logHandler_(message);
    }
}

void KvmClient::OnMessage(const std::string& message) {
    if (message.empty()) {
        LogMessage("Received empty message.");
        return;
    }

    KVMMessage kvmMsg;
    if (!kvmMsg.Unpack(message)) {
        LogMessage("Failed to unpack message.");
        return;
    }

    uint16_t cmdType = kvmMsg.GetCmdType();
    auto it = recvMessageHandlers_.find(cmdType);
    if (it != recvMessageHandlers_.end()) {
        // 调用注册的回调函数
        it->second(kvmMsg.GetBody());
    } else {
        LogMessage("No handler registered for command type: " + std::to_string(cmdType));
    }
}

void KvmClient::ProcessConnectRsp(const std::vector<uint8_t>& body) {
    if (body.size() != 1) {
        LogMessage("Invalid response size.");
        return;
    }

    uint8_t rspCode = body[0];

    {
        std::lock_guard<std::mutex> lock(connectRspMutex_);
        connectRsp_ = (rspCode == 0) ? E_OK : E_ERR; // 根据响应码设置结果
    }
    connectRspCv_.notify_one(); // 通知等待的线程

    switch (rspCode) {
    case 0:
        LogMessage("Connect successfully.");
        break;
    case 1:
        LogMessage("Connect failed.");
        break;
    default:
        LogMessage("Unknown response code. RspCode: " + std::to_string(rspCode));
        break;
    }
}
